import glob
import cv2
import numpy as np
import os
import math
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


# change_here
# ============================================================================
root = "/home/data/323/"
save_path = "/project/train/analysis_dataset/{}"
if not os.path.exists(save_path.format("")):
    os.mkdir(save_path.format(""))
# ============================================================================

# load imgs paths
imgpaths = glob.glob("{root}/*.jpg".format(root=root))
print("len(imgpaths) : ",len(imgpaths))

CLASSES={
    "gas_bottle":0,
}
class DrawLabels:
    def __init__(self,CLASSES):
        self.CLASSES=CLASSES
        self.CLASSES_COLORS={
            self.CLASSES[j]:self.hsv2rgb(i/len(self.CLASSES)*360,1.0,1.0) for i,j in enumerate(self.CLASSES.keys())
        }
        self.font = cv2.FONT_HERSHEY_SIMPLEX
    @staticmethod
    def hsv2rgb(h, s, v):
        h = float(h)
        s = float(s)
        v = float(v)
        h60 = h / 60.0
        h60f = math.floor(h60)
        hi = int(h60f) % 6
        f = h60 - h60f
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        r, g, b = 0, 0, 0
        if hi == 0: r, g, b = v, t, p
        elif hi == 1: r, g, b = q, v, p
        elif hi == 2: r, g, b = p, v, t
        elif hi == 3: r, g, b = p, q, v
        elif hi == 4: r, g, b = t, p, v
        elif hi == 5: r, g, b = v, p, q
        r, g, b = int(r * 255), int(g * 255), int(b * 255)
        return r, g, b
    def draw(self,img,anns):
        # anns=[
        #     [label,x_center,y_center,w,h],
        #     ...,
        #     [label,x_center,y_center,w,h],
        # ]
        img=img.copy()
        for ann in anns:
            label, x, y, w, h = ann
            img = cv2.rectangle(img, (int(x-w/2),int(y-h/2)), (int(x+w/2),int(y+h/2)),self.CLASSES_COLORS[label],2)
            img = cv2.putText(img, str(label), (int(x), int(y)), self.font, 1.2, self.CLASSES_COLORS[label], 2)
        img2 = np.zeros((img.shape[0]+20*(len(self.CLASSES.keys())+1),img.shape[1],img.shape[2]))+128
        img2[:img.shape[0],:img.shape[1],:img.shape[2]]=img
        for i,classname in enumerate(self.CLASSES.keys()):
            img2 = cv2.putText(img2, str(self.CLASSES[classname])+":"+classname, (0, img.shape[0]+(i+1)*20), self.font, 0.8, self.CLASSES_COLORS[CLASSES[classname]], 2)
        return img2
drawlabels = DrawLabels(CLASSES)

def read_imgsize(img,anns):
    return False,list(img.shape)
def show_imgsize(info):
    img_shape = [i["img_shape"] for i in info]
    img_shape = np.array(img_shape)


    plt.figure()
    unq, count = np.unique(img_shape, axis=0, return_counts=True)
    for xy,c in zip(unq,count):
        plt.text(xy[1],xy[0],'%.0f' % c,fontdict={'fontsize':14})
        plt.scatter(xy[1], xy[0], alpha=0.5, s=c*100)
        plt.scatter(xy[1], xy[0], alpha=1, s=1)
    plt.savefig(save_path.format("img_shape_1.jpg"))

    plt.figure()
    plt.scatter(img_shape[:,1], img_shape[:,0])
    plt.savefig(save_path.format("img_shape_2.jpg"))

    plt.figure()
    plt.scatter(img_shape[:,1]*0+1, img_shape[:,1]/img_shape[:,0])
    plt.savefig(save_path.format("img_shape_3_wdivh.jpg"))


def read_ann_shape(img,anns):

    return False,[[i[3],i[4]] for i in anns]
def show_ann_shape(info):

    ann_shape = []
    for i in info:
        ann_shape += i["ann_shape"]
    ann_shape = np.array(ann_shape)

    plt.figure()
    plt.scatter(ann_shape[:,1], ann_shape[:,0])
    plt.savefig(save_path.format("ann_shape.jpg"))

def NONEDEF(*args,**kwargs):
    return None

def read_ann_shape_div_img_shape(img,anns):
    h,w,c = img.shape
    res = []
    tobeshow = False
    for i in anns:
        # if i[3]/w>0.9 and i[4]/h>0.9:
        #     tobeshow = True
        res.append([i[3]/w,i[4]/h])
    return tobeshow,res
def show_ann_shape_div_img_shape(info):
    ann_shape = []
    for i in info:
        ann_shape += i["ann_shape_div_img_shape"]
    ann_shape = np.array(ann_shape)

    plt.figure()
    plt.scatter(ann_shape[:,1], ann_shape[:,0])
    plt.savefig(save_path.format("ann_shape_div_img_shape.jpg"))


infokeys = {
    "img_shape":[read_imgsize,show_imgsize,[]],
    "ann_shape":[read_ann_shape,show_ann_shape,[]],
    "ann_shape_div_img_shape":[read_ann_shape_div_img_shape,show_ann_shape_div_img_shape,["img_shape","ann_shape"]],
}

for key in infokeys.keys():
    for needed_key in infokeys[key][2]:
        if needed_key not in infokeys.keys():
            raise Exception("'{}'' is needed for '{}'".format(needed_key,key))




# read imgs infos
info = []
for imgpath in imgpaths:
    img = cv2.imread(imgpath)
    annspath = imgpath[:-4]+".txt"
    if not os.path.exists(annspath):
        continue
    height,weight = img.shape[:2]

    anns2 = []
    with open(annspath) as f:
        anns = f.readlines()
        for ann in anns:
            label, x, y, w, h = [float(i) for i in ann.split(" ")]
            label = int(label)
            x *= weight
            w *= weight
            y *= height
            h *= height
            anns2.append([label, x, y, w, h])

    dict0 = {}
    for key in infokeys.keys():
        tobeshow,b = infokeys[key][0](img,anns2)
        dict0[key] = b

    if tobeshow:
        img2 = drawlabels.draw(img,anns2)
        cv2.imwrite(save_path+"/tobeshow.jpg",img2)
        1/0
    info.append(dict0)

# show imgs infos
for key in infokeys.keys():
    infokeys[key][1](info)

# IMG SIZE